import json

import torch.cuda
from transformers import BertModel, BertTokenizer, BertConfig

import os

# 获取当前脚本的绝对路径
current_script_path = os.path.abspath(__file__)


class Config:
	def __init__(self):
		# 模型名称
		self.model_name = "bert"
		# 根目录
		# self.root_path = 'F:/GroupFive/TouManFen/'
		# 获取脚本所在的目录
		# self.root_path = os.path.dirname(current_script_path)
		self.root_path = os.path.dirname(os.path.dirname(current_script_path)) + '/'
		# 原始数据路径
		self.train_datapath = self.root_path + '01-data/train.csv'
		self.test_datapath = self.root_path + '01-data/test.csv'
		self.dev_datapath = self.root_path + '01-data/dev.csv'
		# 类别文档
		self.class_doc_path = self.root_path + "/01-data/class.json"
		# 类别字典
		with open(self.class_doc_path, 'r', encoding='utf-8') as f:
			data = json.load(f)['label_to_idx']
		self.id2class_dict = {v: k for k, v in data.items()}

		# 模型保存路径
		self.model_save_path = self.root_path + "04-bert/save_models/bert_classifer_model.pt"
		# 训练设备选择
		self.device = 'cuda' if torch.cuda.is_available() else 'cpu'

		self.bert_path = self.root_path + "/04-bert/bert-base-chinese"  # 预训练Bert模型路径
		self.bert_model = BertModel.from_pretrained(self.bert_path)  # 加载预训练Bert模型
		self.bert_tokenizer = BertTokenizer.from_pretrained(self.bert_path)  # 加载预训练Bert的分词器
		self.bert_config = BertConfig.from_pretrained(self.bert_path)  # Bert模型的配置

		# 超参数设置
		self.num_epoch = 10  # 模型训练轮次
		self.num_classes = len(self.id2class_dict)  # 类别数
		self.batch_size = 64  # 批次数量
		self.pad_size = 75  # token长度
		self.learning_rate = 5e-5  # 学习率

		# 量化模型保存地址
		self.bert_model_quantization_model_path = self.root_path + "04-bert/model/bert_classifer_quantization_model.pt"


if __name__ == '__main__':
	conf = Config()
	# print(conf.device)
	# print(conf.id2class_dict)
	# print(conf.bert_tokenizer)
	# input_size = conf.bert_tokenizer.convert_tokens_to_ids(["你", "好", "吗"])
	# print(input_size)
	# output_tokens = conf.bert_tokenizer.convert_ids_to_tokens([872, 1962, 1408])
	# print(output_tokens)
	# print('============================')
	# print(conf.bert_model)
	# print(conf.bert_config)
	print(conf.model_save_path)
